#include<stdio.h>
#include<math.h>
#include<stdlib.h>
#include<time.h>

#define PI 3.1415926

void DFT(float * input,int Nx,int Nz,float * real,float * imag);
void DFT1(float * input,int Nx,int Nz,float * real,float * imag);
void IDFT(float *real_part,float *imag_part,int Nx,int Nz,float *idft_part);

int main()
{
    int it,fm;
    float t1,t0,N,dt,L2;
    clock_t start,end;
    FILE *fp;
    N=1024*1;
    dt=0.0005;
    t0=0.1;
    fm=25;

    float *ricker=(float *)malloc(sizeof(float)*N);
    int f1 = 100;//   	# 信号的特征频率1
    int f2 = 300;//   # 信号的特征频率2
    int f3 = 500;//   # 信号的特征频率3
    for(int i=0; i<N; i++) {
        //t1=it*dt-t0;
        //ricker[it]=(1.0-2.0*PI*PI*fm*fm*t1*t1)*exp(-(PI*PI*fm*fm*t1*t1));
        float j = i;
        ricker[i] = 10 + 7 * sin(2*PI*f1*(j*1/N)) + 5 * sin(2*PI*f2* (j*1/N)) + 3 * sin(2*PI*f3* (j*1/N));

    }
    fp=fopen("ricker.txt","wb");
    //fwrite(ricker,sizeof(float),N,fp);
    //fclose(fp);
    for(int i = 0; i < N ; i++)
        fprintf(fp,"%f,",ricker[i]);
    fclose(fp);

    float *Real_part=(float *)malloc(sizeof(float)*N);
    float *Imag_part=(float *)malloc(sizeof(float)*N);
    float *Idft_part=(float *)malloc(sizeof(float)*N);
    float *residual =(float *)malloc(sizeof(float)*N);

    start=clock();
    printf("[%s +%d %s] 1*N=%f\n", __FILE__, __LINE__, __func__, 1*N);

    DFT1(ricker,1,N,Real_part,Imag_part);
    DFT(ricker,1,N,Real_part,Imag_part);
    end=clock();
    fprintf(stderr,"DFT time taken: %lf\n",(double)((end-start)/CLOCKS_PER_SEC));

    start=clock();
    //IDFT(Real_part,Imag_part,1,N,Idft_part);
    end=clock();
    fprintf(stderr,"IDFT time taken: %lf\n",(double)((end-start)/CLOCKS_PER_SEC));

    L2=0;
    for(it=0; it<N; it++) {
        residual[it]=ricker[it]-Idft_part[it];
        L2+=residual[it]*residual[it];
    }
    fprintf(stderr,"Residual between DFT with IDFT is %f\n",sqrt(L2));
    fp=fopen("ricker_residual.bin","wb");
    fwrite(residual,sizeof(float),N,fp);
    fclose(fp);

    free(ricker);
    free(Real_part);
    free(Imag_part);
    free(Idft_part);
    free(residual);
    return 0;

}

void DFT(float * input,int Nx,int Nz,float * real,float * imag)
{

    float *amp=(float *)malloc(sizeof(float)*Nx*Nz);
    float R,I;
    float re,im;
    int ix,in,ik;
    for(ix=0; ix<Nx; ix++)
    {
        for(ik=0; ik<Nz; ik++)
        {
            real[ix*Nz+ik]=0;
            imag[ix*Nz+ik]=0;

            R=0;
            I=0;
            for(in=0; in<Nz; in++)
            {
                R=cos(((-2)*PI*ik*in)/Nz);
                I=sin(((-2)*PI*ik*in)/Nz);

                real[ix*Nz+ik] += input[ix*Nz+in]*R;
                imag[ix*Nz+ik] += input[ix*Nz+in]*I;
            }
            re=real[ix*Nz+ik];
            im=imag[ix*Nz+ik];
            amp[ix*Nz+ik]=sqrt(pow(re,2)+pow(im,2));
        }

    }
    printf("\n");
    for(ix=0; ix<Nx; ix++)
    {
        for(ik=0; ik<Nz; ik++)
        {
            printf(" %f", amp[ix*Nz+ik]);
        }

    }
    printf("\n");
    FILE *fp=fopen("ricker_dft_real.txt","wb");
    FILE *fm=fopen("ricker_dft_imag.bin","wb");
    FILE *fr=fopen("ricker_dft_amplitude.txt","wb");

    fwrite(real,sizeof(float),Nx*Nz,fp);
    fwrite(imag,sizeof(float),Nx*Nz,fm);
    //fwrite(amp,sizeof(float),Nx*Nz,fr);
    printf("[%s +%d %s] Nx*Nz=%d\n", __FILE__, __LINE__, __func__, Nx*Nz);

    for(int i = 0; i < Nx*Nz ; i++)
        fprintf(fr,"%f,",amp[i]);

    fclose(fp);
    fclose(fm);
    fclose(fr);

    free(amp);
}
void DFT1(float * input,int Nx,int Nz,float * real,float * imag)
{

    float *amp=(float *)malloc(sizeof(float)*Nx*Nz);
    float R,I;
    float re,im;
    int ix,in,ik;
    for(ix=0; ix<Nx; ix++)
    {
        for(ik=0; ik<Nz; ik++)
        {
            real[ix*Nz+ik]=0;
            imag[ix*Nz+ik]=0;

            R=0;
            I=0;
            for(in=0; in<Nz; in++)
            {
                R=cos(((-2)*PI*ik*in)/Nz);
                I=sin(((-2)*PI*ik*in)/Nz);

                real[ix*Nz+ik] += input[ix*Nz+in]*R;
                imag[ix*Nz+ik] += 0;//input[ix*Nz+in]*I;
            }
            re=real[ix*Nz+ik];
            im=imag[ix*Nz+ik];
            amp[ix*Nz+ik]=sqrt(pow(re,2)+pow(im,2));
        }

    }
    printf("\n");
    for(ix=0; ix<Nx; ix++)
    {
        for(ik=0; ik<Nz; ik++)
        {
            printf(" %f", amp[ix*Nz+ik]);
        }

    }
    printf("\n");
    FILE *fp=fopen("ricker_dft_real.txt","wb");
    FILE *fm=fopen("ricker_dft_imag.bin","wb");
    FILE *fr=fopen("ricker_dft_amplitude1.txt","wb");

    fwrite(real,sizeof(float),Nx*Nz,fp);
    fwrite(imag,sizeof(float),Nx*Nz,fm);
    //fwrite(amp,sizeof(float),Nx*Nz,fr);
    printf("[%s +%d %s] Nx*Nz=%d\n", __FILE__, __LINE__, __func__, Nx*Nz);

    for(int i = 0; i < Nx*Nz ; i++)
        fprintf(fr,"%f,",amp[i]);

    fclose(fp);
    fclose(fm);
    fclose(fr);

    free(amp);
}

void IDFT(float *real_part,float *imag_part,int Nx,int Nz,float *idft_part)
{
    float *real=(float *)malloc(sizeof(float)*Nx*Nz);
    float *imag=(float *)malloc(sizeof(float)*Nx*Nz);

    float R,I;
    int ix,in,ik;
    for(ix=0; ix<Nx; ix++)
    {
        for(in=0; in<Nz; in++)
        {
            real[ix*Nz+in]=0;
            imag[ix*Nz+in]=0;

            R=0;
            I=0;
            for(ik=0; ik<Nz; ik++)
            {
                R=cos((2*PI*ik*in)/Nz);
                I=sin((2*PI*ik*in)/Nz);

                real[ix*Nz+in] += real_part[ix*Nz+ik]*R;
                imag[ix*Nz+in] += imag_part[ix*Nz+ik]*I;
            }
            idft_part[ix*Nz+in]=((float)(real[ix*Nz+in]/Nz)-(float)(imag[ix*Nz+in]/Nz));//a-b !!!!
        }
    }

    FILE *fp=fopen("ricker_idft.bin","wb");
    fwrite(idft_part,sizeof(float),Nx*Nz,fp);

    fclose(fp);
    free(real);
    free(imag);
}